from scapy.all import *
import matplotlib.pyplot as plt

# 具体要求：收集自己本机的网络流量数据（至少1小时）并进行数据显示。
# A、 网络流量大小的时序图，可按每半分钟、每分钟、每五分钟、每十分钟进行分别显示。
# B、 流量协议类型直方图
# C、 可设置过滤条件，显示指定协议数据包、显示时间段数据包、显示长度范围内的数据包


#监听程序,传入监听时间（以秒为单位）
def monitor(time, file_name):
    packets = sniff(timeout = time)  #filter="",iface="any", prn=function, count=N
    wrpcap(file_name, packets)
# monitor(30, 'demo1.pcap')

#time时间戳
class NetFlow:  #网络流量大小时序图
    def __init__(self,file,x_size):
        self.file = file
        self.x_size = x_size
    def x_time(self): #构建x轴步长
        data_x = []
        for x_number in range(int(3600/self.x_size)+1):
            data_x.append(x_number*self.x_size)
        return data_x
    def y_time(self):  #将数据提取出来，同时将时间戳转化为相对时间
        data = rdpcap(self.file)
        data_y = []
        for packet in data:
            data_y.append(int(packet.time - data[0].time))#转化为相对时间
        return data_y
    def drawing(self):
        plt.rcParams['font.sans-serif'] = 'SimHei' #设置rc参数实现显示中文
        print('等待中。。。')
        plt.hist(self.y_time(), bins = self.x_time(), color = 'g')
        plt.xticks(self.x_time()) #控制x轴刻度
        plt.title("流量时序图")
        plt.show()
# b = NetFlow('demo1.pcap',10)
# b.drawing()

class protocol:   #流量协议类型直方图
    def __init__(self, file):
        self.file = file
    def x_protocol(self):
        data_x = ['ipv4', 'icmp', 'igmp', 'tcp', 'udp', 'eigrp', 'ospf', 'arp', 'ipv6', 'other']
        return data_x
    def y_protocol(self):
        data_y = {'ipv4': 0, 'icmp(ipv4)': 0, 'igmp(ipv4)': 0, 'tcp(ipv4)': 0, 'udp(ipv4)': 0, 'eigrp(ipv4)': 0, 'ospf(ipv4)': 0, 'arp': 0, 'ipv6': 0, 'other': 0}
        packet_list = rdpcap(self.file)
        for packet in packet_list:
            if packet.type == 0x0800: #ipv4
                data_y['ipv4'] += 1
                if packet.proto == 1:
                    data_y['icmp(ipv4)'] += 1
                elif packet.proto == 2:
                    data_y['igmp(ipv4)'] += 1
                elif packet.proto == 6:
                    data_y['tcp(ipv4)'] += 1
                elif packet.proto == 17:
                    data_y['udp(ipv4)'] += 1
                elif packet.proto == 88:
                    data_y['eigrp(ipv4)'] += 1
                elif packet.proto == 89:
                    data_y['ospf(ipv4)'] += 1
            elif packet.type == 0x0806: #ARP
                data_y['arp'] += 1
            elif packet.type == 0x86dd: #ipv6
                data_y['ipv6'] += 1
            else:
                data_y['other'] += 1
        data_y_list = []
        for i in data_y:
            data_y_list.append(data_y[i])
        return data_y_list
    def drawing_bar(self):   #修改成柱状图
        plt.figure(figsize = (18, 9))
        plt.rcParams['font.sans-serif'] = 'SimHei' #设置rc参数实现显示中文
        plt.bar(self.x_protocol(), self.y_protocol(), width = 0.35, color = 'y')
        plt.xticks(self.x_protocol()) #定长x轴
        plt.xlabel('协议名称')
        plt.ylabel('数量')
        plt.title('流量协议类型直方图')
        plt.show()
# c = protocol("demo1.pcap")
# c.drawing_bar() 


#可以用packet.show()显示详细数据包内容
class filter:  #按过滤条件显示数据包内容
    def __init__(self, file):
        self.file = file 
        self.packets = rdpcap(self.file)

    def filter_proto(self, proto):
        for packet in self.packets:
            if proto in packet:
                print(packet.getlayer(proto).summary())

    def filter_time(self, begin_time, end_time):
        for packet in self.packets:
            if packet.time-self.packets[0].time >= begin_time and packet.time-self.packets[0].time <= end_time:
                print(packet.summary())

    def filter_size(self, size_floor, size_ceiling):
        for packet in self.packets:
            if len(packet) >= size_floor and len(packet) <= size_ceiling:
                print(packet.summary())  
# d = filter("demo1.pcap")
# d.filter_proto("ARP")
# d.filter_time(10, 30)
# d.filter_size(100, 200)

def main():
    print(' A、 网络流量大小的时序图，可按每半分钟、每分钟、每五分钟、每十分钟进行分别显示.')
    print(' B、 流量协议类型直方图.')
    print(' C、 可设置过滤条件，显示指定协议数据包、显示时间段数据包、显示长度范围内的数据包')
    choose = input('选择查看 A, B or C :')
    check_file = input('选择文件:')
    if choose == 'A':
        choose_a = int(input("选择时间间隔(单位:秒)："))
        netflow = NetFlow(check_file, choose_a)
        netflow.drawing()
    elif choose == 'B':
        proto = protocol(check_file)
        proto.drawing_bar()
    elif choose == 'C':
        filt = filter(check_file)
        print('选择过滤条件：')
        print('A.显示指定协议数据包')
        print('B.显示时间段数据包')
        print('C.显示长度范围内的数据包')
        choose_c = input('选择查看 A, B or C :')
        if choose_c == 'A':
            protocol_name = input('输入查看的协议名(大写):')
            filt.filter_proto(protocol_name)
        elif choose_c == 'B':
            start = int(input('想要查看数据包的起始时间:'))
            end = int(input('想要查看数据包的终止时间:'))
            filt.filter_time(start, end)
        elif choose_c == 'C':
            min = int(input("最小长度:"))
            max = int(input("最大长度:"))
            filt.filter_size(min, max)
        else:
            print('输入错误！')
    else:
        print("输入错误！")

if __name__ == "__main__":
    while(1):
        main()